# import torch
# import torch
# from transformers import BertModel, BertConfig
# # 初始化模型和优化器
#  # 定义BERT模型的配置参数
# config = BertConfig(
#     vocab_size=30522,  # 根据BERT的词汇量设置
#     hidden_size=768,
#     num_hidden_layers=12,
#     num_attention_heads=12,
#     intermediate_size=3072,
#     max_position_embeddings=512
# )

# # 根据定义的配置创建BERT模型
# model = BertModel(config)
# optimizer = torch.optim.Adam(model.parameters(), lr=1e-5)

# # 指定批次大小
# batch_size = 1 # 你需要根据你的GPU容量调整这个值

# # 确保CUDA可用
# if torch.cuda.is_available():
#     model = model.cuda()

# # 初始化记录最大显存使用量
# max_memory_used = 0

# for i in range(3):
#     # 创建随机输入数据并将其转移到GPU
#     inputs = torch.rand([batch_size, 512, 1024], requires_grad=True)
#     if torch.cuda.is_available():
#         inputs = inputs.cuda()
#     # 为BERT生成随机输入ID，确保是长整型
#     input_ids = torch.randint(0, config.vocab_size, (batch_size, 512), dtype=torch.long)
#     # 创建全1的注意力掩码
#     attention_mask = torch.ones_like(input_ids)

#     output = model(input_ids, attention_mask=attention_mask)
    
#     # 提取输出的最后隐藏状态，计算求和
#     logits = output.last_hidden_state
#     loss = logits.sum()
#     loss.backward()
#     optimizer.step()
#     # 清除旧的梯度
#     optimizer.zero_grad()
    
#     # 检查和记录当前步骤后的最大显存使用
#     current_memory = torch.cuda.max_memory_allocated()
#     max_memory_used = max(max_memory_used, current_memory)
    
#     # 可选：在每次迭代后重置最大内存跟踪器
#     torch.cuda.reset_peak_memory_stats()

# # 打印记录的最大显存使用量
# print(f"Maximum GPU memory used: {max_memory_used / (1024**3)} GB")
import torch
from transformers import BertModel, BertConfig
import torch.optim as optim
import time

# 设置设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 定义BERT模型的配置参数
config = BertConfig(
    vocab_size=30522,  # 根据BERT的词汇量设置
    hidden_size=768,
    num_hidden_layers=12,
    num_attention_heads=12,
    intermediate_size=3072,
    max_position_embeddings=512
)

# 根据定义的配置创建BERT模型
model = BertModel(config)
model = model.to(device)

# 设置优化器
optimizer = optim.Adam(model.parameters(), lr=1e-5)

# 指定批次大小
batch_size = 1 # 根据你的GPU容量调整这个值

# 初始化记录最大显存使用量
max_memory_used = 0

time_start = time.time()
for i in range(3):  # 简单示例：运行3个训练步骤
    # 创建随机输入ID
    input_ids = torch.randint(0, config.vocab_size, (batch_size, 512), dtype=torch.long).to(device)
    attention_mask = torch.ones_like(input_ids).to(device)

    # 前向传播
    outputs = model(input_ids=input_ids, attention_mask=attention_mask)
    loss = outputs.last_hidden_state.sum()  # 使用简单的loss函数

    # 反向传播
    loss.backward()

    # 更新模型
    optimizer.step()

    # 检查并记录当前步骤后的最大显存使用
    current_memory = torch.cuda.max_memory_allocated(device)
    max_memory_used = max(max_memory_used, current_memory)
    
    # 重置显存峰值统计
    torch.cuda.reset_peak_memory_stats(device)
time_end = time.time()
# 打印记录的最大显存使用量
print(time_end-time_start)
print(f"Maximum GPU memory used: {max_memory_used / (1024**3):.15f} GB")

